import pickle
import cv2
import face_recognition
import time

import importlib
import sys

if sys.getdefaultencoding() != 'utf-8':
    importlib.reload(sys)
    sys.setdefaultencoding('utf-8')


class FaceRecognition:
    def __init__(self, target_name_count=2, timeout_seconds=60):
        self.target_name_count = target_name_count
        self.recognized_names = set()
        self.ans = set()  # 使用集合而不是列表
        self.known_face_encodings = []
        self.known_face_labels = []
        self.timeout_seconds = timeout_seconds
        self.start_time = time.time()

    def add_face(self, res):
        res.start()
        ans = res.query()
        for i in ans:
            codeings = pickle.loads(i[2])
            self.known_face_encodings.append(codeings)
            self.known_face_labels.append(i[1])

    
    def add_db(self,image_path,label,res,class_,cn_name):
        face_encoding = face_recognition.face_encodings(face_recognition.load_image_file(image_path))
        if not face_encoding:
            print(f"在 {image_path} 中未找到人脸")
            return
        res.start()
        dat={
            "label":label,
            "feature":pickle.dumps(face_encoding[0]),
            "class_info":class_,
            "cn_name":cn_name

        }
        res.insert(**dat)
        res.close()
        


    def run(self):
        # 打开摄像头
        cap = cv2.VideoCapture(0)

        while True:
            # 读取摄像头帧
            ret, frame = cap.read()

            # 转换颜色空间
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            # 检测当前帧中的人脸
            face_locations = face_recognition.face_locations(rgb_frame)
            face_encodings = face_recognition.face_encodings(rgb_frame, face_locations)

            # 存储当前帧的所有人脸标签
            frame_face_labels = set()  # 使用集合而不是列表

            for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
                # 在人脸周围绘制矩形框
                cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)

                # 尝试匹配人脸
                matches = face_recognition.compare_faces(self.known_face_encodings, face_encoding)
                name = "Unknown"

                # 如果有匹配，则使用第一个匹配的标签
                if True in matches:
                    first_match_index = matches.index(True)
                    name = self.known_face_labels[first_match_index]

                # 在矩形框上方显示匹配的标签
                # cv2.putText(frame, name, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
                cv2.putText(frame, name, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2, cv2.LINE_AA)

                # 添加当前人脸的标签到集合
                frame_face_labels.add(name)
                if name != 'Unknown':
                    self.ans.add(name)

            # 更新已识别的名字集合
            self.recognized_names.update(frame_face_labels)

            # 判断退出条件
            if len(self.recognized_names - {"Unknown"}) >= self.target_name_count or time.time() - self.start_time > self.timeout_seconds:
                print(f"已达到目标人脸数量 ({self.target_name_count}) 或超过时间限制 ({self.timeout_seconds} 秒). 退出...")
                break

            # 显示结果
            cv2.imshow('检测窗口', frame)

            # 按 'q' 键退出循环
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        # 释放摄像头并关闭窗口
        cap.release()
        cv2.destroyAllWindows()
        return self.ans

# # 示例用法：
# face_recognition_instance = FaceRecognition(target_name_count=2, timeout_seconds=20)

# # 添加已知人脸
# face_recognition_instance.add_face('obama.jpg', 'obama')
# face_recognition_instance.add_db('biden.jpg', 'biden')

# 运行人脸识别
# face_recognition_instance.run()
